#!/usr/bin/env python3
"""
Drone Navigator - 无人机导航器
集成A*路径规划算法，实现3D路径规划
"""

import rclpy
from rclpy.node import Node
from rclpy.callback_groups import ReentrantCallbackGroup

from geometry_msgs.msg import PoseStamped, Point
from nav_msgs.msg import OccupancyGrid, Path
from std_msgs.msg import Header
from visualization_msgs.msg import MarkerArray, Marker

import numpy as np
import heapq
from typing import List, Tuple, Optional
import time

class Node3D:
    """3D A*算法节点"""
    def __init__(self, x: int, y: int, z: int):
        self.x = x
        self.y = y  
        self.z = z
        self.g = float('inf')  # 从起点到当前节点的实际成本
        self.h = 0.0           # 启发式成本（到终点的估计成本）
        self.f = float('inf')  # 总成本 f = g + h
        self.parent = None     # 父节点
        
    def __lt__(self, other):
        return self.f < other.f
        
    def __eq__(self, other):
        return self.x == other.x and self.y == other.y and self.z == other.z
        
    def __hash__(self):
        return hash((self.x, self.y, self.z))

class DroneNavigator(Node):
    """
    无人机导航器
    功能：
    1. 接收目标点和当前位置
    2. 基于障碍物地图进行3D A*路径规划
    3. 发布规划路径
    4. 动态重规划
    """
    
    def __init__(self):
        super().__init__('drone_navigator')
        
        # 规划参数
        self.declare_parameter('planning_resolution', 0.2)  # 规划分辨率(米)
        self.declare_parameter('max_planning_time', 10.0)   # 最大规划时间(秒)
        self.declare_parameter('replan_distance', 2.0)      # 重规划触发距离(米)
        self.declare_parameter('flight_height', 2.0)        # 默认飞行高度(米)
        self.declare_parameter('height_layers', 10)         # 高度层数
        self.declare_parameter('max_height', 5.0)           # 最大飞行高度(米)
        
        self.planning_resolution = self.get_parameter('planning_resolution').value
        self.max_planning_time = self.get_parameter('max_planning_time').value
        self.replan_distance = self.get_parameter('replan_distance').value
        self.flight_height = self.get_parameter('flight_height').value
        self.height_layers = self.get_parameter('height_layers').value
        self.max_height = self.get_parameter('max_height').value
        
        # 状态变量
        self.current_pose = None
        self.current_goal = None
        self.occupancy_grid = None
        self.last_planned_goal = None
        self.planning_in_progress = False
        
        # 3D占用网格 [高度][y][x]
        self.grid_3d = None
        self.grid_size_x = 0
        self.grid_size_y = 0
        self.grid_origin_x = 0.0
        self.grid_origin_y = 0.0
        self.grid_resolution = 0.1
        
        # 回调组
        self.callback_group = ReentrantCallbackGroup()
        
        # 发布器
        self.path_pub = self.create_publisher(
            Path, '/planned_path', 10)
        self.debug_marker_pub = self.create_publisher(
            MarkerArray, '/planning_debug', 10)
            
        # 订阅器
        self.goal_sub = self.create_subscription(
            PoseStamped, '/goal_pose', self.goal_callback, 10,
            callback_group=self.callback_group)
        self.pose_sub = self.create_subscription(
            PoseStamped, '/current_pose', self.pose_callback, 10,
            callback_group=self.callback_group)
        self.grid_sub = self.create_subscription(
            OccupancyGrid, '/obstacle_grid', self.grid_callback, 10,
            callback_group=self.callback_group)
            
        # 定时器检查重规划
        self.replan_timer = self.create_timer(
            2.0, self.check_replan, callback_group=self.callback_group)
            
        self.get_logger().info("无人机导航器已启动 - A*路径规划")
        
    def pose_callback(self, msg):
        """当前位置回调"""
        self.current_pose = msg
        
    def goal_callback(self, msg):
        """目标点回调"""
        self.get_logger().info(f"收到规划请求: 目标点({msg.pose.position.x:.2f}, {msg.pose.position.y:.2f}, {msg.pose.position.z:.2f})")
        self.current_goal = msg
        self.last_planned_goal = None  # 重置上次规划的目标
        self.plan_path()
        
    def grid_callback(self, msg):
        """占用栅格地图回调"""
        self.occupancy_grid = msg
        self.update_3d_grid()
        
    def update_3d_grid(self):
        """更新3D占用网格"""
        if self.occupancy_grid is None:
            return
            
        # 更新2D网格信息
        self.grid_size_x = self.occupancy_grid.info.width
        self.grid_size_y = self.occupancy_grid.info.height
        self.grid_resolution = self.occupancy_grid.info.resolution
        self.grid_origin_x = self.occupancy_grid.info.origin.position.x
        self.grid_origin_y = self.occupancy_grid.info.origin.position.y
        
        # 创建3D网格
        self.grid_3d = np.zeros((self.height_layers, self.grid_size_y, self.grid_size_x), dtype=bool)
        
        # 将2D占用网格扩展到3D（每个障碍物从地面延伸到最大高度）
        grid_2d = np.array(self.occupancy_grid.data).reshape(self.grid_size_y, self.grid_size_x)
        
        for z in range(self.height_layers):
            for y in range(self.grid_size_y):
                for x in range(self.grid_size_x):
                    # 如果2D网格中有障碍物，在所有高度层都标记为障碍物
                    if grid_2d[y, x] > 50:  # 占用概率 > 50%
                        self.grid_3d[z, y, x] = True
                        
    def check_replan(self):
        """检查是否需要重规划"""
        if (self.current_pose is None or self.current_goal is None or 
            self.last_planned_goal is None or self.planning_in_progress):
            return
            
        # 检查目标是否变化
        goal_distance = self.calculate_distance_3d(
            self.current_goal.pose.position,
            self.last_planned_goal.pose.position
        )
        
        if goal_distance > self.replan_distance:
            self.get_logger().info("目标变化，触发重规划")
            self.plan_path()
            
    def plan_path(self):
        """执行路径规划"""
        if (self.current_pose is None or self.current_goal is None or 
            self.grid_3d is None or self.planning_in_progress):
            return
            
        self.planning_in_progress = True
        start_time = time.time()
        
        try:
            # 转换起点和终点到网格坐标
            start_grid = self.world_to_grid(self.current_pose.pose.position)
            goal_grid = self.world_to_grid(self.current_goal.pose.position)
            
            if start_grid is None or goal_grid is None:
                self.get_logger().error("起点或终点超出网格范围")
                return
                
            self.get_logger().info(f"开始A*规划: 起点{start_grid} -> 终点{goal_grid}")
            
            # 执行A*算法
            path_grid = self.astar_3d(start_grid, goal_grid)
            
            if path_grid:
                # 转换回世界坐标并发布路径
                path_world = self.grid_path_to_world(path_grid)
                self.publish_path(path_world)
                self.last_planned_goal = self.current_goal
                
                planning_time = time.time() - start_time
                self.get_logger().info(f"路径规划成功! 用时{planning_time:.2f}秒, 路径长度{len(path_world)}点")
                
            else:
                self.get_logger().error("路径规划失败 - 无法找到可行路径")
                
        except Exception as e:
            self.get_logger().error(f"路径规划异常: {str(e)}")
            
        finally:
            self.planning_in_progress = False
            
    def astar_3d(self, start: Tuple[int, int, int], goal: Tuple[int, int, int]) -> Optional[List[Tuple[int, int, int]]]:
        """
        3D A*路径规划算法 - 优化版本
        使用更高效的数据结构和算法优化
        """
        start_time = time.time()
        
        # 创建起点和终点节点
        start_node = Node3D(start[0], start[1], start[2])
        start_node.g = 0
        start_node.h = self.heuristic_optimized(start, goal)
        start_node.f = start_node.g + start_node.h
        
        goal_node = Node3D(goal[0], goal[1], goal[2])
        
        # 开放列表使用堆
        open_heap = [start_node]
        heapq.heapify(open_heap)
        
        # 使用字典提高查找效率
        open_set = {(start[0], start[1], start[2]): start_node}
        closed_set = set()
        
        # 预计算邻居偏移量（26连通性）
        neighbor_offsets = self._get_neighbor_offsets()
        
        # 性能统计
        nodes_explored = 0
        max_nodes = 50000  # 限制搜索节点数
        
        while open_heap and nodes_explored < max_nodes:
            # 检查超时
            if time.time() - start_time > self.max_planning_time:
                self.get_logger().warn(f"A*算法超时，已搜索{nodes_explored}个节点")
                break
                
            # 从开放列表中取出f值最小的节点
            current = heapq.heappop(open_heap)
            current_pos = (current.x, current.y, current.z)
            
            # 从开放列表中移除
            if current_pos in open_set:
                del open_set[current_pos]
            
            # 添加到关闭列表
            closed_set.add(current_pos)
            nodes_explored += 1
            
            # 检查是否到达目标
            if current_pos == goal:
                path = self.reconstruct_path(current)
                self.get_logger().info(f"路径找到! 搜索了{nodes_explored}个节点，用时{time.time()-start_time:.2f}秒")
                return path
            
            # 遍历邻居节点（使用预计算的偏移量）
            for offset in neighbor_offsets:
                neighbor_pos = (current.x + offset[0], current.y + offset[1], current.z + offset[2])
                
                # 边界检查
                if not self._is_valid_position(neighbor_pos):
                    continue
                    
                # 检查是否在关闭列表中
                if neighbor_pos in closed_set:
                    continue
                    
                # 检查是否是障碍物
                if self.is_obstacle(neighbor_pos):
                    continue
                
                # 计算移动成本
                move_cost = self.movement_cost_optimized(current_pos, neighbor_pos)
                tentative_g = current.g + move_cost
                
                # 检查是否在开放列表中
                if neighbor_pos in open_set:
                    neighbor = open_set[neighbor_pos]
                    if tentative_g < neighbor.g:
                        # 找到更好的路径
                        neighbor.g = tentative_g
                        neighbor.f = neighbor.g + neighbor.h
                        neighbor.parent = current
                        heapq.heapify(open_heap)  # 重新排序堆
                else:
                    # 创建新节点
                    neighbor = Node3D(neighbor_pos[0], neighbor_pos[1], neighbor_pos[2])
                    neighbor.g = tentative_g
                    neighbor.h = self.heuristic_optimized(neighbor_pos, goal)
                    neighbor.f = neighbor.g + neighbor.h
                    neighbor.parent = current
                    
                    heapq.heappush(open_heap, neighbor)
                    open_set[neighbor_pos] = neighbor
        
        self.get_logger().error(f"未找到路径! 搜索了{nodes_explored}个节点")
        return None
        
    def _get_neighbor_offsets(self):
        """预计算26连通性的邻居偏移量"""
        offsets = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    offsets.append((dx, dy, dz))
        return offsets
        
    def _is_valid_position(self, pos: Tuple[int, int, int]) -> bool:
        """检查位置是否在有效范围内"""
        x, y, z = pos
        return (0 <= x < self.grid_size_x and 
                0 <= y < self.grid_size_y and 
                0 <= z < self.height_layers)
    
    def heuristic_optimized(self, pos1: Tuple[int, int, int], pos2: Tuple[int, int, int]) -> float:
        """优化的启发式函数 - 使用欧几里得距离"""
        dx = pos1[0] - pos2[0]
        dy = pos1[1] - pos2[1]
        dz = pos1[2] - pos2[2]
        return (dx*dx + dy*dy + dz*dz) ** 0.5
        
    def movement_cost_optimized(self, pos1: Tuple[int, int, int], pos2: Tuple[int, int, int]) -> float:
        """优化的移动成本计算"""
        dx = abs(pos1[0] - pos2[0])
        dy = abs(pos1[1] - pos2[1])
        dz = abs(pos1[2] - pos2[2])
        
        # 对角线移动成本
        if dx + dy + dz == 3:  # 3D对角线
            return 1.732  # sqrt(3)
        elif dx + dy + dz == 2:  # 2D对角线
            return 1.414  # sqrt(2)
        else:  # 直线移动
            # Z轴移动成本稍高（鼓励水平移动）
            return 1.2 if dz > 0 else 1.0
        
    def reconstruct_path(self, node: Node3D) -> List[Tuple[int, int, int]]:
        """重构路径 - 优化版本"""
        path = []
        current = node
        while current is not None:
            path.append((current.x, current.y, current.z))
            current = current.parent
        return path[::-1]  # 使用切片反转更高效
        
    def world_to_grid(self, position: Point) -> Optional[Tuple[int, int, int]]:
        """世界坐标转网格坐标 - 优化版本"""
        if self.grid_3d is None:
            return None
            
        # 计算网格坐标
        grid_x = int((position.x - self.grid_origin_x) / self.grid_resolution)
        grid_y = int((position.y - self.grid_origin_y) / self.grid_resolution)
        grid_z = int(position.z / (self.max_height / self.height_layers))
        
        # 边界检查
        if not self._is_valid_position((grid_x, grid_y, grid_z)):
            return None
            
        return (grid_x, grid_y, grid_z)
        
    def grid_to_world(self, grid_pos: Tuple[int, int, int]) -> Point:
        """网格坐标转世界坐标 - 优化版本"""
        x, y, z = grid_pos
        point = Point()
        point.x = self.grid_origin_x + x * self.grid_resolution
        point.y = self.grid_origin_y + y * self.grid_resolution
        point.z = z * (self.max_height / self.height_layers)
        return point
        
    def grid_path_to_world(self, grid_path: List[Tuple[int, int, int]]) -> List[Point]:
        """批量转换网格路径到世界坐标 - 优化版本"""
        world_path = []
        for grid_pos in grid_path:
            world_path.append(self.grid_to_world(grid_pos))
        return world_path
        
    def publish_path(self, path_points: List[Point]):
        """发布路径 - 优化版本，减少内存分配"""
        path_msg = Path()
        path_msg.header.frame_id = "map"
        path_msg.header.stamp = self.get_clock().now().to_msg()
        
        # 预分配列表容量
        path_msg.poses = [None] * len(path_points)
        
        for i, point in enumerate(path_points):
            pose_stamped = PoseStamped()
            pose_stamped.header = path_msg.header
            pose_stamped.pose.position = point
            pose_stamped.pose.orientation.w = 1.0
            path_msg.poses[i] = pose_stamped
            
        self.path_pub.publish(path_msg)
        
    def calculate_distance_3d(self, pos1: Point, pos2: Point) -> float:
        """计算3D距离 - 优化版本"""
        dx = pos1.x - pos2.x
        dy = pos1.y - pos2.y
        dz = pos1.z - pos2.z
        return (dx*dx + dy*dy + dz*dz) ** 0.5

    def is_obstacle(self, pos: Tuple[int, int, int]) -> bool:
        """检查位置是否为障碍物"""
        x, y, z = pos
        
        if (0 <= x < self.grid_size_x and 
            0 <= y < self.grid_size_y and 
            0 <= z < self.height_layers):
            return self.grid_3d[z, y, x]
        return True  # 边界外视为障碍物

def main(args=None):
    rclpy.init(args=args)
    
    drone_navigator = DroneNavigator()
    
    try:
        rclpy.spin(drone_navigator)
    except KeyboardInterrupt:
        pass
    finally:
        drone_navigator.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main() 